// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <unordered_map>
#include <unordered_set>
#include <vector>

#include <tbir/pipeline/node.h>

namespace tbir::runtime {

    class TXSession;

    class Graph {
    public:
        explicit Graph(const std::vector<NodePtr> &outputs);

        explicit Graph(const NodePtr &output_node);

        virtual ~Graph() = default;

    public:
        static std::shared_ptr<Graph> FromGenericList(TXSession *sess, List generic_graph);

        List ToGenericList() const;

    public:
        /**
         *
         * @return
         */
        const std::vector<NodePtr> &get_topo_nodes() const;

        /**
         *
         * @return
         */
        const std::vector<int> &get_input_nodes_index() const;

        /**
         *
         * @return
         */
        std::vector<NodePtr> get_input_nodes() const;

        /**
         *
         * @return
         */
        const std::vector<int> &get_output_nodes_index() const;

        /**
         *
         * @return
         */
        std::vector<NodePtr> get_output_nodes() const;

    private:
        explicit Graph() = default;

        void topo_sort(const std::vector<NodePtr> &outputs);

        void build_input_output(const std::vector<NodePtr> &outputs);

        static void dfs_visit(std::vector<NodePtr> &nodes,
                              std::unordered_set<NodePtr> &visited,
                              const NodePtr &node);

        void add_entry(NodeEntryPtr e) {
            name2entry_->emplace(e->key, e);
        }

    private:
        std::shared_ptr<std::vector<NodePtr>> nodes_;
        std::shared_ptr<std::vector<int>> inputs_;
        std::shared_ptr<std::vector<int>> outputs_;
        std::shared_ptr<std::unordered_map<std::string, NodeEntryPtr>> name2entry_;

        friend class Node;
    };

}  // namespace tbir::runtime
